#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import socket
import threading
import struct
import serial
from datetime import datetime
from typing import Callable, Optional, Dict, Any

from ..server.data_parse import is_modbus_data, parse_modbus_rtu


class TcpClientManager:
    """TCP客户端管理类"""
    
    def __init__(self, message_callback: Optional[Callable] = None):
        """
        初始化TCP客户端管理器
        
        Args:
            message_callback: 消息回调函数，用于向GUI发送日志消息
                             函数签名: callback(message: str, level: str = "INFO", msg_type: str = "normal")
                             msg_type可以是: "normal", "send", "receive", "error"
        """
        self.tcp_server_socket = None
        self.udp_server_socket = None
        self.modbus_tcp_socket = None
        self.modbus_rtu_connection = None  # 串口连接用于Modbus RTU
        self.running = False

        # 回调函数
        self.message_callback = message_callback or self._default_callback
        
    def _default_callback(self, message: str, level: str = "INFO", msg_type: str = "normal"):
        """默认回调函数，打印到控制台"""
        timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
        print(f"[{timestamp}] {level}: {message}")
    
    def connect_tcp_server(self, host: str, port: int) -> bool:
        """
        连接TCP服务器
        
        Args:
            host: 服务器IP地址
            port: 服务器端口
            
        Returns:
            bool: 连接是否成功
        """
        try:
            if self.tcp_server_socket:
                self.tcp_server_socket.close()
            
            self.tcp_server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.tcp_server_socket.connect((host, port))
            
            self.running = True
            self.message_callback(f"已连接到TCP服务器 {host}:{port}")
            
            # 启动接收线程
            threading.Thread(target=self._tcp_receive_thread, daemon=True).start()
            
            return True
            
        except Exception as e:
            self.message_callback(f"连接TCP服务器失败: {e}", "ERROR", "error")
            return False
    
    def disconnect_tcp_server(self):
        """断开TCP服务器连接"""
        try:
            self.running = False
            
            if self.tcp_server_socket:
                self.tcp_server_socket.close()
                self.tcp_server_socket = None
            
            self.message_callback("已断开TCP服务器连接")
            
        except Exception as e:
            self.message_callback(f"断开TCP服务器连接错误: {e}", "ERROR", "error")

    def connect_udp_server(self, host: str, port: int) -> bool:
        """
        连接UDP服务器

        Args:
            host: 服务器IP地址
            port: 服务器端口

        Returns:
            bool: 连接是否成功
        """
        try:
            if self.udp_server_socket:
                self.udp_server_socket.close()

            self.udp_server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            # UDP不需要connect，只需要记录目标地址
            self.udp_target_address = (host, port)

            self.running = True
            self.message_callback(f"已连接到UDP服务器 {host}:{port}")

            # UDP是无连接协议，不需要接收线程
            return True

        except Exception as e:
            self.message_callback(f"连接UDP服务器失败: {e}", "ERROR", "error")
            return False

    def disconnect_udp_server(self):
        """断开UDP服务器连接"""
        try:
            self.running = False

            if self.udp_server_socket:
                self.udp_server_socket.close()
                self.udp_server_socket = None

            self.message_callback("已断开UDP服务器连接")

        except Exception as e:
            self.message_callback(f"断开UDP服务器连接错误: {e}", "ERROR", "error")

    def connect_server(self, protocol: str, host: str, port: int) -> bool:
        """
        统一的服务器连接方法

        Args:
            protocol: 协议类型 ("TCP" 或 "UDP")
            host: 服务器IP地址
            port: 服务器端口

        Returns:
            bool: 连接是否成功
        """
        if protocol == "TCP":
            return self.connect_tcp_server(host, port)
        elif protocol == "UDP":
            return self.connect_udp_server(host, port)
        else:
            self.message_callback(f"不支持的协议: {protocol}", "ERROR", "error")
            return False

    def disconnect_server(self):
        """
        统一的服务器断开方法
        """
        try:
            # 断开TCP连接
            if self.tcp_server_socket:
                self.disconnect_tcp_server()

            # 断开UDP连接
            if self.udp_server_socket:
                self.disconnect_udp_server()

        except Exception as e:
            self.message_callback(f"断开服务器连接错误: {e}", "ERROR", "error")

    def connect_modbus_device(self, host: str, port: int) -> bool:
        """
        连接Modbus TCP设备
        
        Args:
            host: Modbus设备IP地址
            port: Modbus设备端口
            
        Returns:
            bool: 连接是否成功
        """
        try:
            if self.modbus_tcp_socket:
                self.modbus_tcp_socket.close()
            
            self.modbus_tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.modbus_tcp_socket.settimeout(10)
            self.modbus_tcp_socket.connect((host, port))
            
            self.message_callback(f"已连接到Modbus设备 {host}:{port}")
            
            return True
            
        except Exception as e:
            self.message_callback(f"连接Modbus设备失败: {e}", "ERROR", "error")
            return False
    
    def connect_modbus_rtu_device(self, port: str, baudrate: int = 9600,
                                 databits: int = 8, parity: str = 'N', stopbits: int = 1) -> bool:
        """
        连接Modbus RTU设备

        Args:
            port: 串口名称
            baudrate: 波特率
            databits: 数据位
            parity: 校验位 ('N', 'E', 'O')
            stopbits: 停止位

        Returns:
            bool: 连接是否成功
        """
        try:
            if self.modbus_rtu_connection and self.modbus_rtu_connection.is_open:
                self.modbus_rtu_connection.close()

            self.modbus_rtu_connection = serial.Serial(
                port=port,
                baudrate=baudrate,
                bytesize=databits,
                parity=parity,
                stopbits=stopbits,
                timeout=3
            )

            self.message_callback(f"已连接到Modbus RTU设备: {port}@{baudrate},{databits}{parity}{stopbits}")

            return True

        except Exception as e:
            self.message_callback(f"连接Modbus RTU设备失败: {e}", "ERROR", "error")
            return False

    def disconnect_modbus_device(self):
        """断开Modbus设备连接（TCP和RTU）"""
        try:
            # 断开TCP连接
            if self.modbus_tcp_socket:
                self.modbus_tcp_socket.close()
                self.modbus_tcp_socket = None

            # 断开RTU连接
            if self.modbus_rtu_connection and self.modbus_rtu_connection.is_open:
                self.modbus_rtu_connection.close()
                self.modbus_rtu_connection = None

            self.message_callback("已断开Modbus设备连接")

        except Exception as e:
            self.message_callback(f"断开Modbus设备连接错误: {e}", "ERROR", "error")
    
    def _tcp_receive_thread(self):
        """TCP服务器接收线程"""
        while self.running and self.tcp_server_socket:
            try:
                data = self.tcp_server_socket.recv(1024)
                if data:
                    self.message_callback(f"📨 收到TCP服务器数据: {data.hex().upper()}", "INFO", "receive")
                else:
                    self.message_callback("TCP服务器断开连接", "INFO", "receive")
                    break
            except Exception as e:
                if self.running:
                    self.message_callback(f"TCP接收错误: {e}", "ERROR", "error")
                break

    def send_command(self, cmd_str: str) -> bool:
        """
        发送指令或数据（异步执行）

        Args:
            cmd_str: 要发送的指令字符串

        Returns:
            bool: 发送是否成功（立即返回True，实际结果通过回调通知）
        """
        if not cmd_str.strip():
            return False

        # 在后台线程中执行发送操作，避免阻塞GUI
        def send_in_background():
            try:
                # 尝试判断输入内容的类型
                is_modbus_command = self._is_modbus_command(cmd_str)

                # 优先级：Modbus设备 > 服务器连接
                if is_modbus_command:
                    # 如果连接了Modbus RTU设备，使用RTU协议
                    if self.modbus_rtu_connection and self.modbus_rtu_connection.is_open:
                        self._send_as_modbus_rtu(cmd_str)
                    # 如果连接了Modbus TCP设备，使用TCP协议
                    elif self.modbus_tcp_socket:
                        self._send_as_modbus_tcp(cmd_str)
                    # 否则发送到服务器
                    elif self.tcp_server_socket:
                        self._send_as_raw_data_tcp(cmd_str)
                    elif self.udp_server_socket:
                        self._send_as_raw_data_udp(cmd_str)
                    else:
                        self.message_callback("请先连接服务器或Modbus设备", "ERROR", "error")
                        return
                else:
                    # 非Modbus数据，发送到服务器
                    if self.tcp_server_socket:
                        self._send_as_raw_data_tcp(cmd_str)
                    elif self.udp_server_socket:
                        self._send_as_raw_data_udp(cmd_str)
                    else:
                        self.message_callback("请先连接服务器", "ERROR", "error")
                        return

            except Exception as e:
                self.message_callback(f"发送数据失败: {e}", "ERROR", "error")

        # 启动后台线程
        import threading
        send_thread = threading.Thread(target=send_in_background, daemon=True)
        send_thread.start()

        self.message_callback("⏳ 正在发送数据...", "INFO", "send")
        return True
    
    def _is_modbus_command(self, cmd_str: str) -> bool:
        """判断输入是否为Modbus指令格式"""
        try:
            # 移除空格并检查是否为十六进制格式
            hex_str = cmd_str.replace(' ', '').replace('-', '').replace(':', '')
            
            # 检查是否为有效的十六进制字符串
            if not all(c in '0123456789ABCDEFabcdef' for c in hex_str):
                return False
            
            # 检查长度是否合理（至少6个字节：设备地址+功能码+地址+数量）
            if len(hex_str) < 12:  # 6字节 = 12个十六进制字符
                return False
            
            # 尝试解析为字节数据
            byte_data = bytes.fromhex(hex_str)
            
            # 检查是否符合Modbus RTU格式（简单检查功能码）
            if len(byte_data) >= 2:
                function_code = byte_data[1]
                # 常见的Modbus功能码：1,2,3,4,5,6,15,16
                valid_function_codes = [1, 2, 3, 4, 5, 6, 15, 16]
                return function_code in valid_function_codes
            
            return False
        except:
            return False

    def _send_as_raw_data_tcp(self, cmd_str: str) -> bool:
        """作为原始数据发送到TCP服务器"""
        try:
            byte_data, data_type, original_data = self._parse_data_input(cmd_str)
            self.tcp_server_socket.send(byte_data)
            self.message_callback(f"📤 发送到TCP服务器: {byte_data.hex().upper()}", "INFO", "send")
            if data_type == 'text':
                self.message_callback(f"📝 文本内容: {original_data}", "INFO", "send")
            return True
        except Exception as e:
            self.message_callback(f"发送到TCP服务器失败: {e}", "ERROR", "error")
            return False

    def _send_as_raw_data_udp(self, cmd_str: str) -> bool:
        """作为原始数据发送到UDP服务器"""
        try:
            byte_data, data_type, original_data = self._parse_data_input(cmd_str)
            self.udp_server_socket.sendto(byte_data, self.udp_target_address)
            self.message_callback(f"📤 发送到UDP服务器: {byte_data.hex().upper()}", "INFO", "send")
            if data_type == 'text':
                self.message_callback(f"📝 文本内容: {original_data}", "INFO", "send")
            return True
        except Exception as e:
            self.message_callback(f"发送到UDP服务器失败: {e}", "ERROR", "error")
            return False

    def _send_as_modbus_rtu(self, cmd_str: str) -> bool:
        """作为Modbus RTU指令发送到串口设备"""
        try:
            unit_id, function_code, start_addr, quantity, rtu_with_crc = self._parse_modbus_command(cmd_str)
            self.message_callback(f"📤 解析Modbus RTU指令: 从站{unit_id}, 功能码{function_code}, 地址{start_addr}, 数量{quantity}", "INFO", "send")
            self.message_callback(f"📤 发送RTU指令: {rtu_with_crc.hex().upper()}", "INFO", "send")

            # 发送指令到串口
            self.modbus_rtu_connection.write(rtu_with_crc)
            self.message_callback(f"✅ RTU指令发送成功", "INFO", "send")

            # 接收响应
            response = self.modbus_rtu_connection.read(256)
            if response:
                self.message_callback(f"📨 接收RTU响应: {response.hex().upper()}", "INFO", "receive")

                # 解析响应并转发到服务器
                if self.tcp_server_socket:
                    self.tcp_server_socket.send(response)
                    self.message_callback(f"📤 RTU响应已转发到TCP服务器", "INFO", "send")
                elif self.udp_server_socket:
                    self.udp_server_socket.sendto(response, self.udp_target_address)
                    self.message_callback(f"📤 RTU响应已转发到UDP服务器", "INFO", "send")

                # 解析并显示响应数据
                self._parse_rtu_response(response, function_code, unit_id)
            else:
                self.message_callback("⚠️ 未收到RTU响应", "WARNING", "normal")

            return True

        except Exception as e:
            self.message_callback(f"Modbus RTU发送失败: {e}", "ERROR", "error")
            return False

    def _send_as_modbus_tcp(self, cmd_str: str) -> bool:
        """作为Modbus TCP指令发送"""
        try:
            unit_id, function_code, start_addr, quantity, rtu_with_crc = self._parse_modbus_command(cmd_str)
            self.message_callback(f"📤 解析Modbus指令: 从站{unit_id}, 功能码{function_code}, 地址{start_addr}, 数量{quantity}", "INFO", "send")
            self.message_callback(f"📤 RTU格式: {rtu_with_crc.hex().upper()}", "INFO", "send")
            
            # 构建Modbus TCP PDU
            transaction_id = 0x0001
            protocol_id = 0x0000
            length = 6
            
            modbus_tcp_data = struct.pack(
                ">HHHBBHH",
                transaction_id,
                protocol_id,
                length,
                unit_id,
                function_code,
                start_addr,
                quantity
            )
            
            # 发送TCP请求到Modbus设备
            self.modbus_tcp_socket.send(modbus_tcp_data)
            self.message_callback(f"📤 发送Modbus TCP到设备: {modbus_tcp_data.hex().upper()}", "INFO", "send")
            
            # 接收响应
            response = self.modbus_tcp_socket.recv(1024)
            self.message_callback(f"📨 接收Modbus响应: {response.hex().upper()}", "INFO", "receive")
            
            # 转换为RTU格式并发送到TCP服务器
            if response and self.tcp_server_socket:
                rtu_response = self._convert_tcp_to_rtu_response(response)
                if rtu_response:
                    self.tcp_server_socket.send(rtu_response)
                    self.message_callback(f"📤 RTU数据已发送到TCP服务器: {rtu_response.hex().upper()}", "INFO", "send")
                    
                    # 解析并显示数据
                    parsed = self._parse_modbus_tcp_response(response, function_code)
                    if parsed and parsed.get('success'):
                        self.message_callback(f"✅ 采集成功 从站{unit_id}: {parsed['data']}", "INFO", "receive")
                        if 'registers' in parsed:
                            self.message_callback(f"📊 从站{unit_id} 寄存器值: {parsed['registers']}", "INFO", "receive")
            
            return True
            
        except Exception as e:
            self.message_callback(f"Modbus TCP发送失败: {e}", "ERROR", "error")
            return False

    def _parse_data_input(self, data_input: str) -> tuple:
        """解析输入的数据，支持文本和十六进制格式"""
        if data_input.startswith('text:'):
            # 文本数据
            text_data = data_input[5:]  # 去掉 'text:' 前缀
            return text_data.encode('utf-8'), 'text', text_data
        else:
            # 尝试解析为十六进制数据
            try:
                hex_str = data_input.replace(' ', '').replace('-', '').replace(':', '')
                if len(hex_str) % 2 != 0:
                    hex_str = '0' + hex_str
                byte_data = bytes.fromhex(hex_str)
                return byte_data, 'hex', data_input
            except ValueError:
                # 如果不是有效的十六进制，当作文本处理
                return data_input.encode('utf-8'), 'text', data_input

    def _parse_modbus_command(self, cmd_str: str) -> tuple:
        """解析Modbus RTU格式指令"""
        try:
            parts = cmd_str.strip().split()
            if len(parts) < 6:
                raise ValueError("指令格式应为: 单元ID 功能码 起始地址H 起始地址L 数量H 数量L [CRCH CRCL]")

            unit_id = int(parts[0], 16)
            function_code = int(parts[1], 16)
            start_addr = (int(parts[2], 16) << 8) + int(parts[3], 16)
            quantity = (int(parts[4], 16) << 8) + int(parts[5], 16)

            # 构建RTU数据部分(不带CRC)
            rtu_data = bytes([unit_id, function_code]) + struct.pack(">HH", start_addr, quantity)

            # 检查是否已经包含CRC
            if len(parts) >= 8:
                provided_crc = bytes([int(parts[6], 16), int(parts[7], 16)])
                calculated_crc = self._calculate_crc(rtu_data)

                if provided_crc != calculated_crc:
                    self.message_callback(f"警告: 提供的CRC ({provided_crc.hex().upper()}) 与计算的CRC ({calculated_crc.hex().upper()}) 不匹配")

                crc = calculated_crc
            else:
                crc = self._calculate_crc(rtu_data)
                self.message_callback(f"自动计算CRC: {crc.hex().upper()}")

            return unit_id, function_code, start_addr, quantity, rtu_data + crc

        except ValueError as e:
            raise ValueError(f"无效指令: {e}")

    def _calculate_crc(self, data: bytes) -> bytes:
        """计算Modbus RTU CRC校验码"""
        crc = 0xFFFF
        for pos in data:
            crc ^= pos
            for _ in range(8):
                if (crc & 1) != 0:
                    crc >>= 1
                    crc ^= 0xA001
                else:
                    crc >>= 1
        return crc.to_bytes(2, 'little')

    def _convert_tcp_to_rtu_response(self, tcp_response: bytes) -> Optional[bytes]:
        """将Modbus TCP响应转换为RTU格式"""
        if not tcp_response or len(tcp_response) < 8:
            return None

        try:
            # 提取RTU部分（去掉TCP头）
            rtu_data = tcp_response[6:]  # 从unit_id开始的所有数据

            # 计算CRC并添加
            crc = self._calculate_crc(rtu_data)
            rtu_with_crc = rtu_data + crc

            return rtu_with_crc

        except Exception as e:
            self.message_callback(f"TCP转RTU错误: {e}", "ERROR", "error")
            return None

    def _parse_modbus_tcp_response(self, response: bytes, function_code: int) -> Optional[Dict[str, Any]]:
        """解析Modbus TCP响应数据"""
        if not response or len(response) < 8:
            return None

        try:
            # 解析Modbus TCP头
            (trans_id, prot_id, length, unit_id, func_code) = struct.unpack(">HHHBB", response[:8])

            if func_code == function_code:
                byte_count = response[8]
                data = response[9:9+byte_count]

                # 如果是读取寄存器的响应
                if function_code in [3, 4] and byte_count % 2 == 0:
                    registers = []
                    for i in range(0, byte_count, 2):
                        registers.append((data[i] << 8) + data[i+1])
                    return {
                        'success': True,
                        'data': data.hex().upper(),
                        'registers': registers
                    }
                else:
                    return {
                        'success': True,
                        'data': data.hex().upper()
                    }
            else:
                error_code = response[8] if len(response) > 8 else 0
                return {
                    'success': False,
                    'error': f"功能码={func_code}, 错误码={error_code}"
                }
        except Exception as e:
            self.message_callback(f"解析响应错误: {e}", "ERROR", "error")
            return None

    def is_tcp_server_connected(self) -> bool:
        """检查TCP服务器是否已连接"""
        return self.tcp_server_socket is not None

    def is_modbus_device_connected(self) -> bool:
        """检查Modbus设备是否已连接"""
        return self.modbus_tcp_socket is not None
    
    def _send_as_raw_data(self, cmd_str: str) -> bool:
        """作为原始数据发送到TCP服务器"""
        try:
            # 解析输入数据
            byte_data, data_type, original_data = self._parse_data_input(cmd_str)
            
            # 发送到TCP服务器
            self.tcp_server_socket.send(byte_data)
            
            if data_type == 'hex':
                self.message_callback(f"📤 发送十六进制数据到TCP服务器: {original_data}", "INFO", "send")
                self.message_callback(f"   字节数据: {byte_data.hex().upper()}", "INFO", "send")
            else:
                self.message_callback(f"📤 发送文本数据到TCP服务器: {original_data}", "INFO", "send")
                self.message_callback(f"   字节数据: {byte_data.hex().upper()}", "INFO", "send")
            
            return True
            
        except Exception as e:
            self.message_callback(f"原始数据发送失败: {e}", "ERROR", "error")
            return False

    def _parse_rtu_response(self, response: bytes, function_code: int, unit_id: int):
        """解析RTU响应数据"""
        try:
            if is_modbus_data(response):
                parsed = parse_modbus_rtu(response)
                self.message_callback("┌─ Modbus RTU响应解析 ─────────────────", "INFO", "receive")
                self.message_callback(f"│ 从站ID: {unit_id}", "INFO", "receive")
                self.message_callback(f"│ 功能码: {function_code}", "INFO", "receive")

                if "error" in parsed:
                    self.message_callback(f"│ ❌ 解析失败: {parsed['error']}", "ERROR", "receive")
                else:
                    self.message_callback(f"│ CRC校验: {'✅ 正确' if parsed['crc_valid'] else '❌ 错误'}", "INFO", "receive")

                    if "registers" in parsed:
                        self.message_callback(f"│ 寄存器数量: {len(parsed['registers'])}", "INFO", "receive")
                        for i, reg_val in enumerate(parsed['registers']):
                            self.message_callback(f"│   寄存器{i}: {reg_val}", "INFO", "receive")

                self.message_callback("└─────────────────────────────────────", "INFO", "receive")
            else:
                self.message_callback("📄 非标准Modbus响应", "INFO", "receive")

        except Exception as e:
            self.message_callback(f"解析RTU响应失败: {e}", "ERROR", "error")

    def is_tcp_server_connected(self) -> bool:
        """检查是否连接到TCP服务器"""
        return self.tcp_server_socket is not None

    def is_udp_server_connected(self) -> bool:
        """检查是否连接到UDP服务器"""
        return self.udp_server_socket is not None

    def is_modbus_tcp_connected(self) -> bool:
        """检查是否连接到Modbus TCP设备"""
        return self.modbus_tcp_socket is not None

    def is_modbus_rtu_connected(self) -> bool:
        """检查是否连接到Modbus RTU设备"""
        return self.modbus_rtu_connection is not None and self.modbus_rtu_connection.is_open

    def get_connection_status(self) -> dict:
        """获取所有连接状态"""
        return {
            'tcp_server': self.is_tcp_server_connected(),
            'udp_server': self.is_udp_server_connected(),
            'modbus_tcp': self.is_modbus_tcp_connected(),
            'modbus_rtu': self.is_modbus_rtu_connected()
        }
